import os

os.environ["CUDA_VISIBLE_DEVICES"] = '2,3'
from langchain.agents import create_sql_agent
from langchain.agents.agent_toolkits import SQLDatabaseToolkit
from langchain.agents.agent_types import AgentType
from langchain.sql_database import SQLDatabase

from langchain.llms import VLLM
from langchain import OpenAI, SQLDatabase
from langchain_experimental.sql import SQLDatabaseChain
from langchain.chains import create_sql_query_chain
import jsonlines
import traceback
from modelscope import AutoModelForCausalLM, AutoTokenizer, snapshot_download
from modelscope import GenerationConfig
# from sentence_transformers import SentenceTransformer
import chromadb

model_dir = '/datasets/fengjiahao/nlp/TongyiFinance/Tongyi-Finance-14B'
# model_dir = '/datasets/fengjiahao/nlp/TongyiFinance/Tongyi-Finance-14B-Chat/'
# model_dir = '/datasets/fengjiahao/nlp/qwen/Qwen-7B-Chat/'
# model_dir = '/datasets/fengjiahao/nlp/TongyiFinance/Tongyi-Finance-14B-Chat-Int4/'
question_json_path = r'/datasets/fengjiahao/nlp/bs_challenge_financial_14b_dataset/question.json'
answer_path = r'/datasets/fengjiahao/nlp/bs_challenge_financial_14b_dataset/submit_result.jsonl'


client = chromadb.PersistentClient(path="db/pdf")
collection = client.get_collection(name="pdf")
# embedding_model = SentenceTransformer('/datasets/fengjiahao/nlp/m3e-base/')
tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)



content = []
with jsonlines.open(question_json_path, "r") as json_file:
    for obj in json_file.iter(type=dict, skip_invalid=True):
        content.append(obj)
llm = VLLM(
    model=model_dir,
    trust_remote_code=True,  # mandatory for hf models
    temperature=0.2,
    top_p=0.7,
    top_k=15,
    tensor_parallel_size=2, verbose=True
)
# llm = AutoModelForCausalLM.from_pretrained(model_dir, device_map="auto", trust_remote_code=True).eval()
# llm.generation_config = GenerationConfig.from_pretrained(model_dir, trust_remote_code=True)
table_dict = {
    # 'A股': ['A股票日行情表', 'A股股票行业划分表', 'A股股票行情'],
    'A股': [ 'A股股票行情'],
    '基金': ['基金份额持有人结构', '基金债券持仓明细', '基金可转债持仓明细', '基金基本信息', '基金日行情表', '基金股票持仓明细', '基金规模变动表'],
    '港股': ['港股票日行情表'],
}
tips = '''
问题是:请帮我计算，在20210108，中信行业分类划分的一级行业为综合金融行业中，涨跌幅最大股票的股票代码是？涨跌幅是多少？百分数保留两位小数。股票涨跌幅定义为：（收盘价 - 前一日收盘价 / 前一日收盘价）* 100%。请写出查询数据库的SQL(字段名与表需打上双引号):
SELECT "股票代码", ROUND((("收盘价" - "昨收盘") / "昨收盘") * 100, 2) AS "涨跌幅" FROM "A股股票行情" WHERE "一级行业名称" = '综合金融' AND "交易日" = '20210108' ORDER BY "涨跌幅" DESC LIMIT 1
'''

def ask_llm(ori_question):
    if '公司' not in ori_question:
        theme = 'A股'
        if '基金' in ori_question:
            theme = '基金'
        elif '港股' in ori_question or '香港' in ori_question:
            theme = '港股'

        target_table_list = table_dict[theme]
        db = SQLDatabase.from_uri(
            "sqlite:////datasets/fengjiahao/nlp/bs_challenge_financial_14b_dataset/dataset/博金杯比赛数据.db",
            include_tables=target_table_list,
            view_support=True,
            sample_rows_in_table_info=2)
        question = db.table_info + "\n%s\n问题是：%s请写出查询数据库的SQL(字段名与表需打上双引号):" %(tips,ori_question)

        response = llm(question).split('<|endoftext|>')[0]
        print('!!!response1:',response)
        try:
            db_result = db.run(response)
        except Exception as e:
            traceback.print_exc()
            print(e)
            question = db.table_info + "\n问题是:%s为了查询数据库得到问题的答案,请修正如下有语法错误的SQL语句:%s"%(ori_question,response)
            response = llm(question).split('<|endoftext|>')[0]
            print('!!!response2:',response)
            try:
                db_result = db.run(response)
            except Exception as e:
                traceback.print_exc()
                print(e)
                db_result = '数据库中无此数据。'

        try:
            question = "使用SQL(%s)查询数据库得到的结果如下:%s。结合SQL回答问题：%s\n答案：" %( response,db_result , ori_question)
            response = llm(question).split('<|endoftext|>')[0]
        except Exception as e:
            traceback.print_exc()
            print(e)
    else:
        # q_embedding = embedding_model.encode(ori_question).tolist()
        # results = collection.query(query_texts=ori_question, n_results=1, include=["documents"])
        ref = ''
        # for result in results["documents"]:
        #     for i in result:
        #         ref += i

        question = '参考资料如下:%s\n请回答问题:%s\n答案:' % (ref, ori_question)
        response = llm(question).split('<|endoftext|>')[0]
    print("!!!Q:", question)
    print("!!!A:", response)
    return response


for cont in content:
    question = cont['question']
    response = ask_llm(question)
    # print(question)

    cont['answer'] = response
    # 市盈率是最常用来评估股价水平是否合理的指标之一，是指股票价格与每股盈利的比率。...
with jsonlines.open(answer_path, "w") as json_file:
    json_file.write_all(content)
